﻿using Microsoft.Win32;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace iESELibrary.DESCrypto.Vault
{
	public class SysKey : IDisposable
	{
		#region Constants
		private const int TOKEN_ADJUST_PRIVILEGES = 0x00000020;
		private const int SE_PRIVILEGE_ENABLED = 0x00000002;
		private const int TOKEN_QUERY = 0x00000008;

		private const string SE_BACKUP_NAME = "SeBackupPrivilege";
		private const string SE_RESTORE_NAME = "SeRestorePrivilege";

		private const string LOCAL_HIVE_NAME = "localSystemHive";

		private static readonly byte[] KEY_PERMUTATION_MATRIX = new byte[]{
			0x8, 0x5, 0x4, 0x2,
			0xb, 0x9, 0xd, 0x3,
			0x0, 0x6, 0x1, 0xc,
			0xe, 0xa, 0xf, 0x7
		};
		private static int BOOT_KEY_SIZE = 16;
		private static int HIVE_CLASS_SIZE = 4;
		#endregion

		#region Variables
		private bool disposed = false;
		private List<RegKey> hives;
		private string filePath;
		private byte[] bootKey;
		private string bootKeyHex;
		#endregion

		#region Properties
		public byte[] BootKey
		{
			get
			{
				return this.bootKey;
			}
		}
		#endregion

		#region Constructors
		/// <summary>
		/// Creates an instance of Hives class
		/// </summary>
		/// <param name="file">Full absolute SYSTEM file path</param>
		public SysKey(string path)
		{
			this.hives = new List<RegKey>(){
				new RegKey("JD", @"ControlSet001\Control\Lsa\", "Lookup"),
				new RegKey("Skew1", @"ControlSet001\Control\Lsa\", "SkewMatrix"),
				new RegKey("GBG", @"ControlSet001\Control\Lsa\", "GrafBlumGroup"),
				new RegKey("Data", @"ControlSet001\Control\Lsa\", "Pattern")
			};

			this.bootKey = new byte[BOOT_KEY_SIZE];

			this.filePath = path;
			this.MountExport();
			this.DeobfuscateBootKey();
		}
		#endregion

		#region Enums
		private enum RegSAM : int
		{
			QueryValue = 0x0001,
			SetValue = 0x0002,
			CreateSubKey = 0x0004,
			EnumerateSubKeys = 0x0008,
			Notify = 0x0010,
			CreateLink = 0x0020,
			WOW64_32Key = 0x0200,
			WOW64_64Key = 0x0100,
			WOW64_Res = 0x0300,
			Read = 0x00020019,
			Write = 0x00020006,
			Execute = 0x00020019,
			AllAccess = 0x000f003f
		}

		private enum HotKey : uint
		{
			HKEY_CLASSES_ROOT =		0x80000000,
			HKEY_CURRENT_USER =		0x80000001,
			HKEY_LOCAL_MACHINE =	0x80000002,
			HKEY_USERS =			0x80000003
		};
		#endregion

		#region Struct
		[StructLayout(LayoutKind.Sequential)]
		public struct LUID
		{
			public int LowPart;
			public int HighPart;
		}

		[StructLayout(LayoutKind.Sequential)]
		public struct TOKEN_PRIVILEGES
		{
			public LUID Luid;
			public int Attributes;
			public int PrivilegeCount;
		}
		#endregion

		#region Public Methods
		public byte[] GetByteBootKey()
		{
			return Enumerable.Range(0, this.bootKeyHex.Length)
					 .Where(hex => hex % 2 == 0)
					 .Select(hex => Convert.ToByte(this.bootKeyHex.Substring(hex, 2), 16))
					 .ToArray();
		}
		#endregion

		#region Private Imported Methods
		[DllImport("advapi32.dll")]
		private static extern int RegLoadKey(HotKey hKey, string lpSubKey, string lpFile);

		[DllImport("advapi32.dll")]
		private static extern int RegUnLoadKey(HotKey hKey, string lpSubKey);

		[DllImport("advapi32.dll")]
		private static extern int OpenProcessToken(int ProcessHandle, int DesiredAccess, ref int tokenhandle);
	
		[DllImport("kernel32.dll")]
		private static extern int GetCurrentProcess();
		
		[DllImport("advapi32.dll")]
		private static extern int AdjustTokenPrivileges(int tokenhandle, int disableprivs, [MarshalAs(UnmanagedType.Struct)]ref TOKEN_PRIVILEGES Newstate, int bufferlength, int PreivousState, int Returnlength);
		
		[DllImport("advapi32.dll")]
		private static extern int LookupPrivilegeValue(string lpsystemname, string lpname, [MarshalAs(UnmanagedType.Struct)] ref LUID lpLuid);

		[DllImport("advapi32.dll", EntryPoint = "RegQueryInfoKey", CallingConvention = CallingConvention.Winapi, SetLastError = true)]
		extern private static int RegQueryInfoKey(
			UIntPtr hKey,
			StringBuilder lpClass,
			ref uint lpcbClass,
			IntPtr lpReserved,
			IntPtr lpcSubKeys,
			IntPtr lpcbMaxSubKeyLen,
			IntPtr lpcbMaxClassLen,
			IntPtr lpcValues,
			IntPtr lpcbMaxValueNameLen,
			IntPtr lpcbMaxValueLen,
			IntPtr lpcbSecurityDescriptor,
			IntPtr lpftLastWriteTime
		);

		[DllImport("advapi32.dll", CharSet = CharSet.Auto)]
		extern private static int RegOpenKeyEx(
		  HotKey hKey,
		  string subKey,
		  int ulOptions,
		  RegSAM samDesired,
		  out UIntPtr hkResult);

		//[DllImport("advapi32.dll", SetLastError = true)]
		//private static extern int RegLoadAppKey(String hiveFile, out int hKey, RegSAM samDesired, uint options, int reserved);

	
		#endregion

		#region Private Methods
		/// <summary>
		/// Method mounts SYSTEM file and exports SysKey hive
		/// </summary>
		/// <param name="filePath">Full absolute SYSTEM file path</param>
		private void MountExport()
		{
			int rtnVal = -1;
			int token = 0;
			int result = 0;

			LUID BackupLuid = new LUID();
			LUID RestoreLuid = new LUID();

			TOKEN_PRIVILEGES privilegesToken1 = new TOKEN_PRIVILEGES();
			TOKEN_PRIVILEGES privilegesToken2 = new TOKEN_PRIVILEGES();

			result = OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, ref token);
			result = LookupPrivilegeValue(null, SE_BACKUP_NAME, ref BackupLuid);
			result = LookupPrivilegeValue(null, SE_RESTORE_NAME, ref RestoreLuid);

			privilegesToken1.PrivilegeCount = 1;
			privilegesToken1.Attributes = SE_PRIVILEGE_ENABLED;
			privilegesToken1.Luid = RestoreLuid;

			privilegesToken2.PrivilegeCount = 1;
			privilegesToken2.Attributes = SE_PRIVILEGE_ENABLED;
			privilegesToken2.Luid = BackupLuid;

			result = AdjustTokenPrivileges(token, 0, ref privilegesToken1, 1024, 0, 0);
			result = AdjustTokenPrivileges(token, 0, ref privilegesToken2, 1024, 0, 0);

			rtnVal = RegLoadKey(HotKey.HKEY_LOCAL_MACHINE, LOCAL_HIVE_NAME, filePath);

			if(rtnVal != 0)
			{
				switch (rtnVal)
				{
					case 5:
						throw new Exception("Access Denied");
					case 1314:
						throw new Exception("Privileged Access Required (UAC)");
				}
			}
			else
			{
				foreach(RegKey regKey in this.hives)
				{
					UIntPtr hKey;
					regKey.KeyClass = new StringBuilder(HIVE_CLASS_SIZE);
					uint classSize = (uint)HIVE_CLASS_SIZE * 2 + 1;

					string subKey = String.Format(@"{0}\{1}\{2}", LOCAL_HIVE_NAME, regKey.RelativePath, regKey.Name);
					rtnVal = RegOpenKeyEx(HotKey.HKEY_LOCAL_MACHINE, subKey, 0, RegSAM.Read, out hKey);
					rtnVal = RegQueryInfoKey(hKey, regKey.KeyClass, ref classSize, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero, IntPtr.Zero);

					using(RegistryKey baseKey =  Registry.LocalMachine.OpenSubKey(subKey) )
					{
						regKey.Data = (byte[])baseKey.GetValue(regKey.Value);
					}
				}
				rtnVal = RegUnLoadKey(HotKey.HKEY_LOCAL_MACHINE, LOCAL_HIVE_NAME);
			}	
		}

		/// <summary>
		/// Method deobfuscates bootkey from RegKey class products
		/// </summary>
		private void DeobfuscateBootKey()
		{
			foreach(RegKey regKey in this.hives)
			{
				this.bootKeyHex += regKey.KeyClass.ToString();
			}
			this.bootKey = GetByteBootKey();
			//for (int i = 0; i < 4; i++ )
			//{
			//	Array.Copy(UnicodeEncoding.Unicode.GetBytes(this.hives[i].KeyClass.ToString()), 0, this.bootKey, i * this.hives[i].KeyClassSize, this.hives[i].KeyClassSize);
			//}

			Permutate(ref this.bootKey, KEY_PERMUTATION_MATRIX);
			this.bootKeyHex = BitConverter.ToString(this.bootKey);
		}

		/// <summary>
		/// Method permutates UTF-16 elements against permutation matrix
		/// </summary>
		/// <param name="array">UTF-16 array as StringBuilder that will be permutated</param>
		/// <param name="matrix">Permutation 2-D matrix</param>
		private void Permutate(ref byte[] array, byte[] matrix)
		{
			byte[] buffer = new byte[BOOT_KEY_SIZE];

			for (int i = 0; i < (array.Length > matrix.Length ? matrix.Length : array.Length); i++)
			{
				buffer[i] = array[matrix[i]];
			}
			array = buffer;
		}
		#endregion

		#region IDisposable Members
		public void Dispose()
		{
			Dispose(true);
			GC.SuppressFinalize(this);
		}

		protected virtual void Dispose(bool disposing)
		{
			if (!this.disposed)
			{
				if (disposing)
				{
					RegUnLoadKey(HotKey.HKEY_LOCAL_MACHINE, LOCAL_HIVE_NAME);
					this.hives.Clear();
					this.hives = null;
					this.bootKey = null;
				}

				this.disposed = true;
			}
		}
		#endregion
	}
}